#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Harbor配额检查告警工具 - Python版本

功能：
1. 通过Harbor API检查各项目存储配额使用情况
2. 当使用率超过阈值时发送邮件告警
3. 支持调试模式和详细日志记录

使用环境变量配置:
- HARBOR_URL: Harbor服务地址
- HARBOR_USERNAME: Harbor管理员用户名
- HARBOR_PASSWORD: Harbor管理员密码
- SMTP_SERVER: SMTP服务器地址
- SMTP_PORT: SMTP端口
- MAIL_FROM: 发件人邮箱
- MAIL_TO: 收件人邮箱
- SMTP_USERNAME: SMTP用户名
- SMTP_PASSWORD: SMTP密码
- THRESHOLD: 告警阈值(百分比)
"""

import logging
import requests
import base64
import datetime
import smtplib
from email.mime.text import MIMEText
from typing import Dict, Optional, List
from config_loader import config


class HarborAPIClient:
    """Harbor API客户端"""
    def __init__(self):
        self.base_url = str(config.get("HARBOR_URL"))
        self.username = str(config.get("HARBOR_USERNAME"))
        self.password = str(config.get("HARBOR_PASSWORD"))
        self.session = requests.Session()
        self._setup_auth()
        
    def _setup_auth(self) -> None:
        """设置API认证"""
        auth_str = f"{self.username}:{self.password}"
        self.auth_header = {
            "Authorization": "Basic " + base64.b64encode(
                auth_str.encode()
            ).decode()
        }
        
    def get_quotas(self) -> List[Dict]:
        """获取所有项目配额数据
        
        Returns:
            List[Dict]: 包含所有项目配额数据的列表
        Raises:
            requests.exceptions.RequestException: 当API请求失败时抛出
        """
        url = f"https://{self.base_url}/api/v2.0/quotas"
        try:
            response = self.session.get(
                url,
                headers=self.auth_header,
                verify=False  # 忽略SSL证书验证
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            logging.error(f"获取Harbor配额数据失败: {str(e)}")
            raise


class QuotaChecker:
    """配额检查逻辑"""
    
    def __init__(self, threshold: float = 80.0):
        self.threshold = threshold
        
    def check_quota(self, quota_data: Dict) -> Optional[Dict]:
        """检查单个项目配额
        
        Args:
            quota_data: 单个项目的配额数据
            
        Returns:
            Dict: 包含告警信息的字典(如果超过阈值)
            None: 如果未超过阈值
        """
        try:
            project_name = quota_data["ref"]["name"]
            storage_used = quota_data["used"]["storage"]
            storage_limit = quota_data["hard"]["storage"]
            
            if storage_limit == -1 or storage_limit is None:
                return None
                
            usage_percent = (storage_used / storage_limit) * 100
            
            if usage_percent > self.threshold:
                return {
                    "project": project_name,
                    "usage": round(usage_percent, 2),
                    "used": storage_used,
                    "limit": storage_limit
                }
            return None
        except (KeyError, TypeError) as e:
            logging.error(f"处理配额数据出错: {str(e)}")
            return None


class EmailNotifier:
    """邮件通知类"""
    
    def __init__(self):
        self.smtp_server = str(config.get("SMTP_SERVER"))
        self.smtp_port = int(config.get("SMTP_PORT", 465))
        self.mail_from = str(config.get("MAIL_FROM"))
        self.mail_to = str(config.get("MAIL_TO"))
        self.smtp_username = str(config.get("SMTP_USERNAME"))
        self.smtp_password = str(config.get("SMTP_PASSWORD"))
        self.use_tls = bool(config.get("SMTP_USE_TLS", True))
        
    def send_alert(self, project: str, usage: float) -> bool:
        """发送配额告警邮件
        
        Args:
            project: 项目名称
            usage: 存储使用率
            
        Returns:
            bool: 邮件是否发送成功
        """
        subject = f"Harbor配额告警: 项目 {project} 存储使用率达到 {usage}%"
        body = f"""项目名称: {project}
存储使用率: {usage}%
时间: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}

请及时处理"""
        
        try:
            if self.use_tls:
                server = smtplib.SMTP_SSL(self.smtp_server, self.smtp_port)
            else:
                server = smtplib.SMTP(self.smtp_server, self.smtp_port)
                if self.use_tls:
                    server.starttls()
                    
            server.login(self.smtp_username, self.smtp_password)
            
            msg = MIMEText(body, 'plain', 'utf-8')
            msg['Subject'] = subject
            msg['From'] = self.mail_from
            msg['To'] = self.mail_to
            
            server.send_message(msg)
            server.quit()
            return True
            
        except Exception as e:
            logging.error(f"发送邮件失败: {str(e)}")
            return False


class HarborQuotaChecker:
    """Harbor配额检查主类"""
    def __init__(self, debug: bool = False):
        self.debug = debug
        self.threshold = float(config.get("THRESHOLD", 80))
        self.logger = self._setup_logging()
        self.api_client = HarborAPIClient()
        self.notifier = EmailNotifier()
        
    def _setup_logging(self) -> logging.Logger:
        """配置日志记录
        
        Returns:
            logging.Logger: 配置好的日志记录器
        """
        logger = logging.getLogger(__name__)
        logger.setLevel(logging.DEBUG if self.debug else logging.INFO)
        
        formatter = logging.Formatter(
            '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        
        # 控制台处理器
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(formatter)
        logger.addHandler(console_handler)
        
        # 文件处理器
        file_handler = logging.FileHandler('harbor_quota_checker.log')
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)
        
        return logger
        
    def check_quotas(self) -> None:
        """检查所有项目配额"""
        self.logger.info("开始检查Harbor存储配额")
        
        try:
            quotas = self.api_client.get_quotas()
            self.logger.debug(f"获取到 {len(quotas)} 个项目的配额数据")
            
            quota_checker = QuotaChecker(self.threshold)
            alerts = []
            
            for quota in quotas:
                alert = quota_checker.check_quota(quota)
                if alert:
                    alerts.append(alert)
                    self.logger.warning(
                        f"项目 {alert['project']} 存储使用率 {alert['usage']}% 超过阈值")
                    
            if alerts:
                for alert in alerts:
                    success = self.notifier.send_alert(
                        alert['project'], alert['usage'])
                    if not success:
                        self.logger.error(
                            f"发送项目 {alert['project']} 告警邮件失败")
            
            self.logger.info(f"检查完成，共发现 {len(alerts)} 个告警")
            
        except Exception as e:
            self.logger.error(f"配额检查过程中发生错误: {str(e)}")
            raise


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--debug", help="启用调试模式", action="store_true")
    args = parser.parse_args()
    
    checker = HarborQuotaChecker(debug=True if args.debug else False)
    checker.check_quotas()